package com.wt.admin.service.ai.impl.agents;

import cn.hutool.json.JSONUtil;
import com.wt.admin.domain.dto.ai.AgentsInfoDTO;
import com.wt.admin.domain.entity.ai.AgentsInfoEntity;
import com.wt.admin.domain.entity.ai.MCPEntity;
import com.wt.admin.domain.entity.ai.ModelConfigEntity;
import com.wt.admin.service.ai.MCPService;
import com.wt.admin.service.ai.ModelConfigService;
import jakarta.annotation.Resource;
import lombok.Data;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
public class AgentsManager {

    @Resource
    AgentsBuilderStrategy agentsBuilderStrategy;
    @Resource
    ModelConfigService modelConfigService;
    @Resource
    MCPService mcpService;


    Map<Integer, Agent> map = new ConcurrentHashMap<>();

    public void create(AgentsInfoDTO agents) {
        String jsonStr = JSONUtil.toJsonStr(agents.getConfig());
        AgentsInfoEntity.Config config = JSONUtil.toBean(jsonStr, AgentsInfoEntity.Config.class);
        ModelConfigEntity vo = modelConfigService.findById(config.getModel().getChatModel());
        List<MCPEntity> mcpEntities = mcpService.findByIds(config.getTools());
        AgentsManager.Agent builder = agentsBuilderStrategy.builder(agents, vo,mcpEntities);
        map.put(agents.getId(), builder);
    }

    public AgentsManager.Agent get(Integer id){
        return map.get(id);
    }

    public void remove(Integer id){
        map.remove(id);
    }


    @Data
    public static class Agent{
        ChatClient chatClient;
        List<ToolCallbackProvider> toolCallbackProvider;

        public Agent(ChatClient chatClient, List<ToolCallbackProvider> toolCallbackProvider) {
            this.chatClient = chatClient;
            this.toolCallbackProvider = toolCallbackProvider;
        }
    }

}
